{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9-final"
    },
    "colab": {
      "name": "Robotic Manipulation - Lecture 5 Breakout.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EgiF12Hf1Dhs",
        "colab_type": "text"
      },
      "source": [
        "This notebook is meant for our lecture 5 breakout session.  I've put together some questions to guide your [here](https://itempool.com/MIT-Robotic-Manipulation/c/uMW2Pmz_fPa).  The points are not real, and will not contribute to your grade!\n",
        "\n",
        "# Notebook setup\n",
        "\n",
        "The following cell will:\n",
        "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
        "- launch a server for our 3D visualizer (MeshCat) that will be used for the remainder of this notebook.\n",
        "\n",
        "You will need to rerun this cell if you restart the kernel, but it should be fast because the machine will already have drake installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eeMrMI0-1Dhu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import importlib\n",
        "import sys\n",
        "from urllib.request import urlretrieve\n",
        "\n",
        "# Install drake.\n",
        "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
        "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
        "                \"setup_manipulation_colab.py\")\n",
        "    from setup_manipulation_colab import setup_manipulation\n",
        "    setup_manipulation(manipulation_sha='9eef22b4c0c66268debe69adef12ce140a55ca05', drake_version='20200918', drake_build='nightly')\n",
        "\n",
        "# Install pyngrok.\n",
        "server_args = []\n",
        "if 'google.colab' in sys.modules:\n",
        "  server_args = ['--ngrok_http_tunnel']\n",
        "\n",
        "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
        "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
        "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
        "\n",
        "# Determine if this notebook is currently running as a notebook or a unit test.\n",
        "from IPython import get_ipython\n",
        "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
        "\n",
        "# Let's do all of our imports here, too.\n",
        "import numpy as np\n",
        "from functools import partial\n",
        "from IPython.display import display\n",
        "from ipywidgets import FloatSlider, Textarea, Layout\n",
        "\n",
        "from pydrake.all import (DiagramBuilder, AddMultibodyPlantSceneGraph, RigidTransform, ConnectMeshcatVisualizer,\n",
        "                         JacobianWrtVariable, Parser, FixedOffsetFrame,\n",
        "                         PiecewisePolynomial, MathematicalProgram, LinearConstraint, Solve)\n",
        "from manipulation.meshcat_utils import plot_mathematical_program\n",
        "from manipulation.utils import FindResource\n",
        "from manipulation.scenarios import AddTwoLinkIiwa\n",
        "\n",
        "# This one is specific to this notebook, but I'm putting it in the header to make it less distracting.\n",
        "def Visualizer(MakeMathematicalProgram):\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
        "    twolink = AddTwoLinkIiwa(plant, q0=[0.0, 0.0])\n",
        "    hand = plant.GetFrameByName(\"iiwa_link_ee\")\n",
        "    plant.Finalize()\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(builder,\n",
        "                                  scene_graph,\n",
        "                                  zmq_url=zmq_url)\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "    plant_context = plant.GetMyContextFromRoot(context)\n",
        "\n",
        "    meshcat.vis.delete()\n",
        "    #meshcat.vis[\"/Background\"].set_property('visible',False)\n",
        "    meshcat.vis[\"/Background\"].set_property('top_color', [0, 0, 0])\n",
        "    meshcat.vis[\"/Background\"].set_property('bottom_color', [0, 0, 0])\n",
        "    meshcat.vis[\"/Grid\"].set_property('visible',False)\n",
        "    meshcat.load()\n",
        "\n",
        "    jacobian = Textarea(value=\"\", description=\"J_G: \", layout={'width':'200pm','height':'50px'}, style={'description_width':'initial'})\n",
        "    display(jacobian)\n",
        "\n",
        "    X, Y = np.meshgrid(np.linspace(-5, 5, 35), np.linspace(-5, 5, 31))\n",
        "\n",
        "    def visualize(q, v_Gdesired=[1.0, 0.0], t=None):\n",
        "        if t:\n",
        "            context.SetTime(t)\n",
        "        plant.SetPositions(plant_context, q)\n",
        "        diagram.Publish(context)\n",
        "\n",
        "        J_G = plant.CalcJacobianTranslationalVelocity(plant_context, JacobianWrtVariable.kQDot, hand, [0,0,0], plant.world_frame(), plant.world_frame())\n",
        "        J_G = J_G[[0,2],:]  # Ignore Y.\n",
        "        jacobian.value = np.array2string(J_G, formatter={'float': lambda x: \"{:5.2f}\".format(x)})\n",
        "\n",
        "        prog = MakeMathematicalProgram(q, J_G, v_Gdesired)\n",
        "        result = Solve(prog)\n",
        "        v = meshcat.vis[\"QP\"]\n",
        "        plot_mathematical_program(v, prog, X, Y, result=result)\n",
        "        # TODO: Add set_object to meshcat.Animation\n",
        "        if False: # meshcat._is_recording:\n",
        "            with meshcat._animation.at_frame(\n",
        "                    v, meshcat._recording_frame_num) as m:\n",
        "                plot_mathematical_program(m, prog, X, Y, result=result)\n",
        "\n",
        "    return visualize, meshcat"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import meshcat.geometry as g\n",
        "import meshcat.transformations as tf\n",
        "from manipulation.meshcat_utils import plot_surface\n",
        "import pydrake.all\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hSeD3PBotJUU",
        "colab_type": "text"
      },
      "source": [
        "# Differential Inverse Kinematics as a Quadratic Program\n",
        "\n",
        "## Define your mathematical program here.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6F_KZoc1QDkY",
        "colab_type": "code",
        "tags": [],
        "colab": {}
      },
      "source": [
        "def MakeMathematicalProgram(q, J_G, v_Gdesired):\n",
        "    prog = MathematicalProgram()\n",
        "    v = prog.NewContinuousVariables(2, 'v')\n",
        "    v_max = 3.0 \n",
        "\n",
        "    error = J_G.dot(v) - v_Gdesired\n",
        "    prog.AddCost(error.dot(error))\n",
        "    prog.AddBoundingBoxConstraint(-v_max, v_max, v)\n",
        "\n",
        "    return prog"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4i4xIYtmrIa-",
        "colab_type": "text"
      },
      "source": [
        "## Visualize a particular joint angle"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "CM9vi0TtrIa_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "visualize, meshcat = Visualizer(MakeMathematicalProgram)\n",
        "\n",
        "q = [-np.pi/2.0+0.5, 1.0]\n",
        "v_Gdesired = [0.5, 0.]\n",
        "visualize(q, v_Gdesired)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ttJ90nfLrIbD",
        "colab_type": "text"
      },
      "source": [
        "## Animated joint trajectory (passing through the singularity)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "ugsWyLNOrIbE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import time\n",
        "\n",
        "visualize, meshcat = Visualizer(MakeMathematicalProgram)\n",
        "\n",
        "v_Gdesired = [1.0, 0.0]\n",
        "T = 2.\n",
        "q = PiecewisePolynomial.FirstOrderHold(\n",
        "    [0, T, 2*T], np.array([[-np.pi / 2.0 + 1., -np.pi / 2.0 - 1., -np.pi / 2.0 + 1.], \n",
        "                           [2., -2., 2]]))\n",
        "\n",
        "nx = 35\n",
        "ny = 31\n",
        "X, Y = np.meshgrid(np.linspace(-5, 5, nx), np.linspace(-5, 5, ny))\n",
        "D = np.vstack((X.reshape(1,-1), Y.reshape(1,-1)))\n",
        "for i in range(2):\n",
        "    for t in np.linspace(0, 2*T, num=100):\n",
        "        visualize(q.value(t), v_Gdesired, t=t)\n",
        "        time.sleep(0.05)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GPayCPx7rIbI",
        "colab_type": "text"
      },
      "source": [
        "## Trajectory slider"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "KbPqBtd2rIbJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "visualize, meshcat = Visualizer(MakeMathematicalProgram)\n",
        "\n",
        "v_Gdesired = [1.0, 0.0]\n",
        "T = 2.\n",
        "qtraj = PiecewisePolynomial.FirstOrderHold(\n",
        "    [0, T], np.array([[-np.pi / 2.0 + 1., -np.pi / 2.0 - 1.], [2., -2.]]))\n",
        "visualize(qtraj.value(0), v_Gdesired)\n",
        "\n",
        "def _t_callback(change):\n",
        "    visualize(qtraj.value(change.new), v_Gdesired)\n",
        "\n",
        "slider = FloatSlider(value=0, min=0, max=T, step=0.05, continuous_update=True, description=\"t\", layout=Layout(width=\"'100'\"))\n",
        "slider.observe(_t_callback, names='value')\n",
        "display(slider)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DNhK1DCOrIbO",
        "colab_type": "text"
      },
      "source": [
        "## Joint Sliders"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "Kn4igrtmrIbO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "visualize, meshcat = Visualizer(MakeMathematicalProgram)\n",
        "\n",
        "q = [-np.pi/2.0 + 0.5, 1.0]\n",
        "v_Gdeised = [1.0, 0.0]\n",
        "visualize(q, v_Gdesired)\n",
        "\n",
        "def _q_callback(change, index):\n",
        "    q[index] = change.new\n",
        "    visualize(q, v_Gdesired)\n",
        "def _vG_callback(change, index):\n",
        "    v_Gdesired[index] = change.new\n",
        "    visualize(q, v_Gdesired)\n",
        "\n",
        "slider = FloatSlider(value=q[0], min=-np.pi, max=np.pi, step=0.1, continuous_update=True, description=\"q0\", layout=Layout(width=\"'100'\"))\n",
        "slider.observe(partial(_q_callback, index=0), names='value')\n",
        "display(slider)\n",
        "\n",
        "slider = FloatSlider(value=q[1], min=-np.pi, max=np.pi, step=0.1, continuous_update=True, description=\"q1\", layout=Layout(width=\"'100'\"))\n",
        "slider.observe(partial(_q_callback, index=1), names='value')\n",
        "display(slider)\n",
        "\n",
        "slider = FloatSlider(value=v_Gdesired[0], min=-4, max=4, step=0.1, continuous_update=True, description=\"v_G_W0\", layout=Layout(width=\"'100'\"))\n",
        "slider.observe(partial(_vG_callback, index=0), names='value')\n",
        "display(slider)\n",
        "\n",
        "slider = FloatSlider(value=v_Gdesired[1], min=-4, max=4, step=0.1, continuous_update=True, description=\"v_G_W1\", layout=Layout(width=\"'100'\"))\n",
        "slider.observe(partial(_vG_callback, index=1), names='value')\n",
        "display(slider)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YamdzS1mrIbT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}